{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# MNIST Training using PyTorch and Step Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "1. [Background](#Background)\n",
    "1. [Setup](#Setup)\n",
    "1. [Data](#Data)\n",
    "1. [Train](#Train)\n",
    "\n",
    "---\n",
    "\n",
    "## Background\n",
    "\n",
    "MNIST is a widely used dataset for handwritten digit classification. It consists of 70,000 labeled 28x28 pixel grayscale images of hand-written digits. The dataset is split into 60,000 training images and 10,000 test images. There are 10 classes (one for each of the 10 digits). This tutorial will show how to train and test an MNIST model on SageMaker using PyTorch.\n",
    "\n",
    "For more information about PyTorch in SageMaker, please visit [sagemaker-pytorch-containers](https://github.com/aws/sagemaker-pytorch-containers) and [sagemaker-python-sdk](https://github.com/aws/sagemaker-python-sdk) github repositories.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting stepfunctions\n",
      "Requirement not upgraded as not directly required: sagemaker>=1.42.8 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from stepfunctions) (1.45.0.dev0)\n",
      "Requirement not upgraded as not directly required: pyyaml in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from stepfunctions) (3.12)\n",
      "Requirement not upgraded as not directly required: boto3>=1.9.213 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from stepfunctions) (1.10.19)\n",
      "Requirement not upgraded as not directly required: protobuf>=3.1 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from sagemaker>=1.42.8->stepfunctions) (3.5.2)\n",
      "Requirement not upgraded as not directly required: requests<2.21,>=2.20.0 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from sagemaker>=1.42.8->stepfunctions) (2.20.0)\n",
      "Requirement not upgraded as not directly required: smdebug-rulesconfig==0.1.2 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from sagemaker>=1.42.8->stepfunctions) (0.1.2)\n",
      "Requirement not upgraded as not directly required: scipy>=0.19.0 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from sagemaker>=1.42.8->stepfunctions) (1.1.0)\n",
      "Requirement not upgraded as not directly required: protobuf3-to-dict>=0.1.5 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from sagemaker>=1.42.8->stepfunctions) (0.1.5)\n",
      "Requirement not upgraded as not directly required: numpy>=1.9.0 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from sagemaker>=1.42.8->stepfunctions) (1.15.4)\n",
      "Requirement not upgraded as not directly required: s3transfer<0.3.0,>=0.2.0 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from boto3>=1.9.213->stepfunctions) (0.2.1)\n",
      "Requirement not upgraded as not directly required: jmespath<1.0.0,>=0.7.1 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from boto3>=1.9.213->stepfunctions) (0.9.4)\n",
      "Requirement not upgraded as not directly required: botocore<1.14.0,>=1.13.19 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from boto3>=1.9.213->stepfunctions) (1.13.19)\n",
      "Requirement not upgraded as not directly required: six>=1.9 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from protobuf>=3.1->sagemaker>=1.42.8->stepfunctions) (1.11.0)\n",
      "Requirement not upgraded as not directly required: setuptools in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from protobuf>=3.1->sagemaker>=1.42.8->stepfunctions) (39.1.0)\n",
      "Requirement not upgraded as not directly required: certifi>=2017.4.17 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from requests<2.21,>=2.20.0->sagemaker>=1.42.8->stepfunctions) (2019.9.11)\n",
      "Requirement not upgraded as not directly required: idna<2.8,>=2.5 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from requests<2.21,>=2.20.0->sagemaker>=1.42.8->stepfunctions) (2.6)\n",
      "Requirement not upgraded as not directly required: urllib3<1.25,>=1.21.1 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from requests<2.21,>=2.20.0->sagemaker>=1.42.8->stepfunctions) (1.23)\n",
      "Requirement not upgraded as not directly required: chardet<3.1.0,>=3.0.2 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from requests<2.21,>=2.20.0->sagemaker>=1.42.8->stepfunctions) (3.0.4)\n",
      "Requirement not upgraded as not directly required: docutils<0.16,>=0.10 in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from botocore<1.14.0,>=1.13.19->boto3>=1.9.213->stepfunctions) (0.14)\n",
      "Requirement not upgraded as not directly required: python-dateutil<2.8.1,>=2.1; python_version >= \"2.7\" in /home/ec2-user/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages (from botocore<1.14.0,>=1.13.19->boto3>=1.9.213->stepfunctions) (2.7.3)\n",
      "\u001b[31mfastai 1.0.59 requires nvidia-ml-py3, which is not installed.\u001b[0m\n",
      "Installing collected packages: stepfunctions\n",
      "Successfully installed stepfunctions-1.0.0.3\n",
      "\u001b[33mYou are using pip version 10.0.1, however version 19.3.1 is available.\n",
      "You should consider upgrading via the 'pip install --upgrade pip' command.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "!{sys.executable} -m pip install --upgrade stepfunctions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Package                            Version        \r\n",
      "---------------------------------- ---------------\r\n",
      "alabaster                          0.7.10         \r\n",
      "anaconda-client                    1.6.14         \r\n",
      "anaconda-project                   0.8.2          \r\n",
      "asn1crypto                         0.24.0         \r\n",
      "astroid                            1.6.3          \r\n",
      "astropy                            3.0.2          \r\n",
      "attrs                              18.1.0         \r\n",
      "autovizwidget                      0.13.1         \r\n",
      "awscli                             1.16.283       \r\n",
      "Babel                              2.5.3          \r\n",
      "backcall                           0.1.0          \r\n",
      "backports.shutil-get-terminal-size 1.0.0          \r\n",
      "bcrypt                             3.1.7          \r\n",
      "beautifulsoup4                     4.6.0          \r\n",
      "bitarray                           0.8.1          \r\n",
      "bkcharts                           0.2            \r\n",
      "blaze                              0.11.3         \r\n",
      "bleach                             2.1.3          \r\n",
      "blis                               0.2.4          \r\n",
      "bokeh                              1.0.4          \r\n",
      "boto                               2.48.0         \r\n",
      "boto3                              1.10.19        \r\n",
      "botocore                           1.13.19        \r\n",
      "Bottleneck                         1.2.1          \r\n",
      "cached-property                    1.5.1          \r\n",
      "certifi                            2019.9.11      \r\n",
      "cffi                               1.11.5         \r\n",
      "chardet                            3.0.4          \r\n",
      "click                              6.7            \r\n",
      "cloudpickle                        0.5.3          \r\n",
      "clyent                             1.2.2          \r\n",
      "colorama                           0.3.9          \r\n",
      "contextlib2                        0.5.5          \r\n",
      "cryptography                       2.8            \r\n",
      "cycler                             0.10.0         \r\n",
      "cymem                              2.0.2          \r\n",
      "Cython                             0.28.2         \r\n",
      "cytoolz                            0.9.0.1        \r\n",
      "dask                               0.17.5         \r\n",
      "dataclasses                        0.6            \r\n",
      "datashape                          0.5.4          \r\n",
      "decorator                          4.3.0          \r\n",
      "defusedxml                         0.6.0          \r\n",
      "distributed                        1.21.8         \r\n",
      "docker                             3.7.3          \r\n",
      "docker-compose                     1.24.1         \r\n",
      "docker-pycreds                     0.4.0          \r\n",
      "dockerpty                          0.4.1          \r\n",
      "docopt                             0.6.2          \r\n",
      "docutils                           0.14           \r\n",
      "entrypoints                        0.2.3          \r\n",
      "environment-kernels                1.1.1          \r\n",
      "et-xmlfile                         1.0.1          \r\n",
      "fabric                             2.5.0          \r\n",
      "fastai                             1.0.59         \r\n",
      "fastcache                          1.0.2          \r\n",
      "fastprogress                       0.1.21         \r\n",
      "filelock                           3.0.4          \r\n",
      "Flask                              1.0.2          \r\n",
      "Flask-Cors                         3.0.4          \r\n",
      "gevent                             1.3.0          \r\n",
      "glob2                              0.6            \r\n",
      "gmpy2                              2.0.8          \r\n",
      "greenlet                           0.4.13         \r\n",
      "h5py                               2.8.0          \r\n",
      "hdijupyterutils                    0.13.1         \r\n",
      "heapdict                           1.0.0          \r\n",
      "html5lib                           1.0.1          \r\n",
      "idna                               2.6            \r\n",
      "imageio                            2.3.0          \r\n",
      "imagesize                          1.0.0          \r\n",
      "invoke                             1.3.0          \r\n",
      "ipykernel                          4.8.2          \r\n",
      "ipyparallel                        6.2.2          \r\n",
      "ipython                            6.4.0          \r\n",
      "ipython-genutils                   0.2.0          \r\n",
      "ipywidgets                         7.4.0          \r\n",
      "isort                              4.3.4          \r\n",
      "itsdangerous                       0.24           \r\n",
      "jdcal                              1.4            \r\n",
      "jedi                               0.12.0         \r\n",
      "Jinja2                             2.10           \r\n",
      "jmespath                           0.9.4          \r\n",
      "jsonschema                         2.6.0          \r\n",
      "jupyter                            1.0.0          \r\n",
      "jupyter-client                     5.2.3          \r\n",
      "jupyter-console                    5.2.0          \r\n",
      "jupyter-core                       4.4.0          \r\n",
      "jupyterlab                         0.32.1         \r\n",
      "jupyterlab-launcher                0.10.5         \r\n",
      "kiwisolver                         1.0.1          \r\n",
      "lazy-object-proxy                  1.3.1          \r\n",
      "llvmlite                           0.23.1         \r\n",
      "locket                             0.2.0          \r\n",
      "lxml                               4.2.1          \r\n",
      "MarkupSafe                         1.0            \r\n",
      "matplotlib                         3.0.3          \r\n",
      "mccabe                             0.6.1          \r\n",
      "mistune                            0.8.3          \r\n",
      "mkl-fft                            1.0.6          \r\n",
      "mkl-random                         1.0.1          \r\n",
      "mock                               3.0.5          \r\n",
      "more-itertools                     4.1.0          \r\n",
      "mpmath                             1.0.0          \r\n",
      "msgpack                            0.6.0          \r\n",
      "msgpack-python                     0.5.6          \r\n",
      "multipledispatch                   0.5.0          \r\n",
      "murmurhash                         1.0.2          \r\n",
      "nb-conda                           2.2.1          \r\n",
      "nb-conda-kernels                   2.2.2          \r\n",
      "nbconvert                          5.4.1          \r\n",
      "nbformat                           4.4.0          \r\n",
      "networkx                           2.1            \r\n",
      "nltk                               3.3            \r\n",
      "nose                               1.3.7          \r\n",
      "notebook                           5.5.0          \r\n",
      "numba                              0.38.0         \r\n",
      "numexpr                            2.6.5          \r\n",
      "numpy                              1.15.4         \r\n",
      "numpydoc                           0.8.0          \r\n",
      "odo                                0.5.1          \r\n",
      "olefile                            0.45.1         \r\n",
      "onnx                               1.5.0          \r\n",
      "openpyxl                           2.5.3          \r\n",
      "packaging                          17.1           \r\n",
      "pandas                             0.24.2         \r\n",
      "pandocfilters                      1.4.2          \r\n",
      "paramiko                           2.6.0          \r\n",
      "parso                              0.2.0          \r\n",
      "partd                              0.3.8          \r\n",
      "path.py                            11.0.1         \r\n",
      "pathlib2                           2.3.2          \r\n",
      "patsy                              0.5.0          \r\n",
      "pep8                               1.7.1          \r\n",
      "pexpect                            4.5.0          \r\n",
      "pickleshare                        0.7.4          \r\n",
      "Pillow                             5.2.0          \r\n",
      "pip                                10.0.1         \r\n",
      "pkginfo                            1.4.2          \r\n",
      "plac                               0.9.6          \r\n",
      "plotly                             4.2.1          \r\n",
      "pluggy                             0.6.0          \r\n",
      "ply                                3.11           \r\n",
      "preshed                            2.0.1          \r\n",
      "prompt-toolkit                     1.0.15         \r\n",
      "protobuf                           3.5.2          \r\n",
      "protobuf3-to-dict                  0.1.5          \r\n",
      "psutil                             5.4.5          \r\n",
      "psycopg2                           2.7.5          \r\n",
      "ptyprocess                         0.5.2          \r\n",
      "py                                 1.5.3          \r\n",
      "py4j                               0.10.7         \r\n",
      "pyasn1                             0.4.8          \r\n",
      "pycodestyle                        2.4.0          \r\n",
      "pycosat                            0.6.3          \r\n",
      "pycparser                          2.18           \r\n",
      "pycrypto                           2.6.1          \r\n",
      "pycurl                             7.43.0.1       \r\n",
      "pyflakes                           1.6.0          \r\n",
      "pygal                              2.4.0          \r\n",
      "Pygments                           2.2.0          \r\n",
      "pykerberos                         1.2.1          \r\n",
      "pylint                             1.8.4          \r\n",
      "PyNaCl                             1.3.0          \r\n",
      "pyodbc                             4.0.23         \r\n",
      "pyOpenSSL                          18.0.0         \r\n",
      "pyparsing                          2.2.0          \r\n",
      "PySocks                            1.6.8          \r\n",
      "pyspark                            2.3.2          \r\n",
      "pytest                             3.5.1          \r\n",
      "pytest-arraydiff                   0.2            \r\n",
      "pytest-astropy                     0.3.0          \r\n",
      "pytest-doctestplus                 0.1.3          \r\n",
      "pytest-openfiles                   0.3.0          \r\n",
      "pytest-remotedata                  0.2.1          \r\n",
      "python-dateutil                    2.7.3          \r\n",
      "pytz                               2018.4         \r\n",
      "PyWavelets                         0.5.2          \r\n",
      "PyYAML                             3.12           \r\n",
      "pyzmq                              17.0.0         \r\n",
      "QtAwesome                          0.4.4          \r\n",
      "qtconsole                          4.3.1          \r\n",
      "QtPy                               1.4.1          \r\n",
      "requests                           2.20.0         \r\n",
      "requests-kerberos                  0.12.0         \r\n",
      "retrying                           1.3.3          \r\n",
      "rope                               0.10.7         \r\n",
      "rsa                                3.4.2          \r\n",
      "ruamel-yaml                        0.15.35        \r\n",
      "s3fs                               0.1.5          \r\n",
      "s3transfer                         0.2.1          \r\n",
      "sagemaker                          1.45.0.dev0    \r\n",
      "sagemaker-pyspark                  1.2.6          \r\n",
      "scikit-image                       0.13.1         \r\n",
      "scikit-learn                       0.20.3         \r\n",
      "scipy                              1.1.0          \r\n",
      "seaborn                            0.8.1          \r\n",
      "Send2Trash                         1.5.0          \r\n",
      "setuptools                         39.1.0         \r\n",
      "simplegeneric                      0.8.1          \r\n",
      "singledispatch                     3.4.0.3        \r\n",
      "six                                1.11.0         \r\n",
      "smdebug-rulesconfig                0.1.2          \r\n",
      "snowballstemmer                    1.2.1          \r\n",
      "sortedcollections                  0.6.1          \r\n",
      "sortedcontainers                   1.5.10         \r\n",
      "spacy                              2.1.8          \r\n",
      "sparkmagic                         0.12.5         \r\n",
      "Sphinx                             1.7.4          \r\n",
      "sphinxcontrib-websupport           1.0.1          \r\n",
      "spyder                             3.2.8          \r\n",
      "SQLAlchemy                         1.2.11         \r\n",
      "srsly                              0.1.0          \r\n",
      "statsmodels                        0.9.0          \r\n",
      "stepfunctions                      1.0.0.3        \r\n",
      "sympy                              1.1.1          \r\n",
      "tables                             3.4.3          \r\n",
      "tblib                              1.3.2          \r\n",
      "terminado                          0.8.1          \r\n",
      "testpath                           0.3.1          \r\n",
      "texttable                          0.9.1          \r\n",
      "thinc                              7.0.8          \r\n",
      "toolz                              0.9.0          \r\n",
      "torch                              1.2.0          \r\n",
      "torchvision                        0.4.0a0+6b959ee\r\n",
      "tornado                            5.0.2          \r\n",
      "tqdm                               4.36.1         \r\n",
      "traitlets                          4.3.2          \r\n",
      "typing                             3.6.4          \r\n",
      "typing-extensions                  3.7.4.1        \r\n",
      "unicodecsv                         0.14.1         \r\n",
      "urllib3                            1.23           \r\n",
      "wasabi                             0.2.2          \r\n",
      "wcwidth                            0.1.7          \r\n",
      "webencodings                       0.5.1          \r\n",
      "websocket-client                   0.56.0         \r\n",
      "Werkzeug                           0.14.1         \r\n",
      "wheel                              0.31.1         \r\n",
      "widgetsnbextension                 3.4.2          \r\n",
      "wrapt                              1.10.11        \r\n",
      "xlrd                               1.1.0          \r\n",
      "XlsxWriter                         1.0.4          \r\n",
      "xlwt                               1.3.0          \r\n",
      "zict                               0.1.3          \r\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mYou are using pip version 10.0.1, however version 19.3.1 is available.\r\n",
      "You should consider upgrading via the 'pip install --upgrade pip' command.\u001b[0m\r\n"
     ]
    }
   ],
   "source": [
    "!pip list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "### Add a policy to your SageMaker role in IAM\n",
    "\n",
    "**If you are running this notebook on an Amazon SageMaker notebook instance**, the IAM role assumed by your notebook instance needs permission to create and run workflows in AWS Step Functions. To provide this permission to the role, do the following.\n",
    "\n",
    "1. Open the Amazon [SageMaker console](https://console.aws.amazon.com/sagemaker/). \n",
    "2. Select **Notebook instances** and choose the name of your notebook instance\n",
    "3. Under **Permissions and encryption** select the role ARN to view the role on the IAM console\n",
    "4. Choose **Attach policies** and search for `AWSStepFunctionsFullAccess`.\n",
    "5. Select the check box next to `AWSStepFunctionsFullAccess` and choose **Attach policy**\n",
    "\n",
    "If you are running this notebook in a local environment, the SDK will use your configured AWS CLI configuration. For more information, see [Configuring the AWS CLI](https://docs.aws.amazon.com/cli/latest/userguide/cli-chap-configure.html).\n",
    "\n",
    "Next, create an execution role in IAM for Step Functions. \n",
    "\n",
    "### Create an execution role for Step Functions\n",
    "\n",
    "You need an execution role so that you can create and execute workflows in Step Functions.\n",
    "\n",
    "1. Go to the [IAM console](https://console.aws.amazon.com/iam/)\n",
    "2. Select **Roles** and then **Create role**.\n",
    "3. Under **Choose the service that will use this role** select **Step Functions**\n",
    "4. Choose **Next** until you can enter a **Role name**\n",
    "5. Enter a name such as `StepFunctionsWorkflowExecutionRole` and then select **Create role**\n",
    "\n",
    "\n",
    "Attach a policy to the role you created. The following steps attach a policy that provides full access to Step Functions, however as a good practice you should only provide access to the resources you need.  \n",
    "\n",
    "1. Under the **Permissions** tab, click **Add inline policy**\n",
    "2. Enter the following in the **JSON** tab\n",
    "\n",
    "```json\n",
    "{\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"sagemaker:CreateTransformJob\",\n",
    "                \"sagemaker:DescribeTransformJob\",\n",
    "                \"sagemaker:StopTransformJob\",\n",
    "                \"sagemaker:CreateTrainingJob\",\n",
    "                \"sagemaker:DescribeTrainingJob\",\n",
    "                \"sagemaker:StopTrainingJob\",\n",
    "                \"sagemaker:CreateHyperParameterTuningJob\",\n",
    "                \"sagemaker:DescribeHyperParameterTuningJob\",\n",
    "                \"sagemaker:StopHyperParameterTuningJob\",\n",
    "                \"sagemaker:CreateModel\",\n",
    "                \"sagemaker:CreateEndpointConfig\",\n",
    "                \"sagemaker:CreateEndpoint\",\n",
    "                \"sagemaker:DeleteEndpointConfig\",\n",
    "                \"sagemaker:DeleteEndpoint\",\n",
    "                \"sagemaker:UpdateEndpoint\",\n",
    "                \"sagemaker:ListTags\",\n",
    "                \"lambda:InvokeFunction\",\n",
    "                \"sqs:SendMessage\",\n",
    "                \"sns:Publish\",\n",
    "                \"ecs:RunTask\",\n",
    "                \"ecs:StopTask\",\n",
    "                \"ecs:DescribeTasks\",\n",
    "                \"dynamodb:GetItem\",\n",
    "                \"dynamodb:PutItem\",\n",
    "                \"dynamodb:UpdateItem\",\n",
    "                \"dynamodb:DeleteItem\",\n",
    "                \"batch:SubmitJob\",\n",
    "                \"batch:DescribeJobs\",\n",
    "                \"batch:TerminateJob\",\n",
    "                \"glue:StartJobRun\",\n",
    "                \"glue:GetJobRun\",\n",
    "                \"glue:GetJobRuns\",\n",
    "                \"glue:BatchStopJobRun\"\n",
    "            ],\n",
    "            \"Resource\": \"*\"\n",
    "        },\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"iam:PassRole\"\n",
    "            ],\n",
    "            \"Resource\": \"*\",\n",
    "            \"Condition\": {\n",
    "                \"StringEquals\": {\n",
    "                    \"iam:PassedToService\": \"sagemaker.amazonaws.com\"\n",
    "                }\n",
    "            }\n",
    "        },\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"events:PutTargets\",\n",
    "                \"events:PutRule\",\n",
    "                \"events:DescribeRule\"\n",
    "            ],\n",
    "            \"Resource\": [\n",
    "                \"arn:aws:events:*:*:rule/StepFunctionsGetEventsForSageMakerTrainingJobsRule\",\n",
    "                \"arn:aws:events:*:*:rule/StepFunctionsGetEventsForSageMakerTransformJobsRule\",\n",
    "                \"arn:aws:events:*:*:rule/StepFunctionsGetEventsForSageMakerTuningJobsRule\",\n",
    "                \"arn:aws:events:*:*:rule/StepFunctionsGetEventsForECSTaskRule\",\n",
    "                \"arn:aws:events:*:*:rule/StepFunctionsGetEventsForBatchJobsRule\"\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "```\n",
    "\n",
    "3. Choose **Review policy** and give the policy a name such as `StepFunctionsWorkflowExecutionPolicy`\n",
    "4. Choose **Create policy**. You will be redirected to the details page for the role.\n",
    "5. Copy the **Role ARN** at the top of the **Summary**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import the required modules \n",
    "\n",
    "Now import the required modules from the Step Functions SDK and AWS SageMaker, configure an S3 bucket, and get the AWS SageMaker execution role."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "import stepfunctions\n",
    "import logging\n",
    "\n",
    "from stepfunctions.template.pipeline import TrainingPipeline\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "stepfunctions.set_stream_logger(level=logging.INFO)\n",
    "\n",
    "bucket = sagemaker_session.default_bucket()\n",
    "prefix = 'sagemaker/DEMO-pytorch-mnist'\n",
    "\n",
    "# SageMaker Execution Role\n",
    "# You can use sagemaker.get_execution_role() if running inside sagemaker's notebook instance\n",
    "sagemaker_execution_role = sagemaker.get_execution_role() #Replace with ARN if not in an AWS SageMaker notebook\n",
    "\n",
    "# paste the StepFunctionsWorkflowExecutionRole ARN from above\n",
    "workflow_execution_role = \"arn:aws:iam::806570384721:role/StepFunctionsExecutionRole\" "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "### Getting the data\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 16384/9912422 [00:00<01:14, 133218.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to data/MNIST/raw/train-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "9920512it [00:00, 40503847.62it/s]                           \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting data/MNIST/raw/train-images-idx3-ubyte.gz to data/MNIST/raw\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "32768it [00:00, 879086.07it/s]\n",
      "1654784it [00:00, 12503228.46it/s]                           "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to data/MNIST/raw/train-labels-idx1-ubyte.gz\n",
      "Extracting data/MNIST/raw/train-labels-idx1-ubyte.gz to data/MNIST/raw\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to data/MNIST/raw/t10k-images-idx3-ubyte.gz\n",
      "Extracting data/MNIST/raw/t10k-images-idx3-ubyte.gz to data/MNIST/raw\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "8192it [00:00, 390447.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to data/MNIST/raw/t10k-labels-idx1-ubyte.gz\n",
      "Extracting data/MNIST/raw/t10k-labels-idx1-ubyte.gz to data/MNIST/raw\n",
      "Processing...\n",
      "Done!\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Dataset MNIST\n",
       "    Number of datapoints: 60000\n",
       "    Root location: data\n",
       "    Split: Train\n",
       "    StandardTransform\n",
       "Transform: Compose(\n",
       "               ToTensor()\n",
       "               Normalize(mean=(0.1307,), std=(0.3081,))\n",
       "           )"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torchvision import datasets, transforms\n",
    "\n",
    "datasets.MNIST('data', download=True, transform=transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.1307,), (0.3081,))\n",
    "]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input spec (in this case, just an S3 path): s3://sagemaker-us-east-1-806570384721/sagemaker/DEMO-pytorch-mnist\n"
     ]
    }
   ],
   "source": [
    "inputs = sagemaker_session.upload_data(path='data', bucket=bucket, key_prefix=prefix)\n",
    "print('input spec (in this case, just an S3 path): {}'.format(inputs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Uploading the data to S3\n",
    "We are going to use the `sagemaker.Session.upload_data` function to upload our datasets to an S3 location. The return value inputs identifies the location -- we will use later when we start the training job.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train\n",
    "### Training script\n",
    "The `mnist.py` script provides all the code we need for training and hosting a SageMaker model (`model_fn` function to load a model).\n",
    "The training script is very similar to a training script you might run outside of SageMaker, but you can access useful properties about the training environment through various environment variables, such as:\n",
    "\n",
    "* `SM_MODEL_DIR`: A string representing the path to the directory to write model artifacts to.\n",
    "  These artifacts are uploaded to S3 for model hosting.\n",
    "* `SM_NUM_GPUS`: The number of gpus available in the current container.\n",
    "* `SM_CURRENT_HOST`: The name of the current container on the container network.\n",
    "* `SM_HOSTS`: JSON encoded list containing all the hosts .\n",
    "\n",
    "Supposing one input channel, 'training', was used in the call to the PyTorch estimator's `fit()` method, the following will be set, following the format `SM_CHANNEL_[channel_name]`:\n",
    "\n",
    "* `SM_CHANNEL_TRAINING`: A string representing the path to the directory containing data in the 'training' channel.\n",
    "\n",
    "For more information about training environment variables, please visit [SageMaker Containers](https://github.com/aws/sagemaker-containers).\n",
    "\n",
    "A typical training script loads data from the input channels, configures training with hyperparameters, trains a model, and saves a model to `model_dir` so that it can be hosted later. Hyperparameters are passed to your script as arguments and can be retrieved with an `argparse.ArgumentParser` instance.\n",
    "\n",
    "Because the SageMaker imports the training script, you should put your training code in a main guard (``if __name__=='__main__':``) if you are using the same script to host your model as we do in this example, so that SageMaker does not inadvertently run your training code at the wrong point in execution.\n",
    "\n",
    "For example, the script run by this notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36margparse\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mjson\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mlogging\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mos\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36msagemaker_containers\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36msys\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtorch\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtorch.distributed\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mdist\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtorch.nn\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mnn\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtorch.nn.functional\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mF\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtorch.optim\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36moptim\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtorch.utils.data\u001b[39;49;00m\r\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mtorch.utils.data.distributed\u001b[39;49;00m\r\n",
      "\u001b[34mfrom\u001b[39;49;00m \u001b[04m\u001b[36mtorchvision\u001b[39;49;00m \u001b[34mimport\u001b[39;49;00m datasets, transforms\r\n",
      "\r\n",
      "logger = logging.getLogger(\u001b[31m__name__\u001b[39;49;00m)\r\n",
      "logger.setLevel(logging.DEBUG)\r\n",
      "logger.addHandler(logging.StreamHandler(sys.stdout))\r\n",
      "\r\n",
      "\r\n",
      "\u001b[37m# Based on https://github.com/pytorch/examples/blob/master/mnist/main.py\u001b[39;49;00m\r\n",
      "\u001b[34mclass\u001b[39;49;00m \u001b[04m\u001b[32mNet\u001b[39;49;00m(nn.Module):\r\n",
      "    \u001b[34mdef\u001b[39;49;00m \u001b[32m__init__\u001b[39;49;00m(\u001b[36mself\u001b[39;49;00m):\r\n",
      "        \u001b[36msuper\u001b[39;49;00m(Net, \u001b[36mself\u001b[39;49;00m).\u001b[32m__init__\u001b[39;49;00m()\r\n",
      "        \u001b[36mself\u001b[39;49;00m.conv1 = nn.Conv2d(\u001b[34m1\u001b[39;49;00m, \u001b[34m10\u001b[39;49;00m, kernel_size=\u001b[34m5\u001b[39;49;00m)\r\n",
      "        \u001b[36mself\u001b[39;49;00m.conv2 = nn.Conv2d(\u001b[34m10\u001b[39;49;00m, \u001b[34m20\u001b[39;49;00m, kernel_size=\u001b[34m5\u001b[39;49;00m)\r\n",
      "        \u001b[36mself\u001b[39;49;00m.conv2_drop = nn.Dropout2d()\r\n",
      "        \u001b[36mself\u001b[39;49;00m.fc1 = nn.Linear(\u001b[34m320\u001b[39;49;00m, \u001b[34m50\u001b[39;49;00m)\r\n",
      "        \u001b[36mself\u001b[39;49;00m.fc2 = nn.Linear(\u001b[34m50\u001b[39;49;00m, \u001b[34m10\u001b[39;49;00m)\r\n",
      "\r\n",
      "    \u001b[34mdef\u001b[39;49;00m \u001b[32mforward\u001b[39;49;00m(\u001b[36mself\u001b[39;49;00m, x):\r\n",
      "        x = F.relu(F.max_pool2d(\u001b[36mself\u001b[39;49;00m.conv1(x), \u001b[34m2\u001b[39;49;00m))\r\n",
      "        x = F.relu(F.max_pool2d(\u001b[36mself\u001b[39;49;00m.conv2_drop(\u001b[36mself\u001b[39;49;00m.conv2(x)), \u001b[34m2\u001b[39;49;00m))\r\n",
      "        x = x.view(-\u001b[34m1\u001b[39;49;00m, \u001b[34m320\u001b[39;49;00m)\r\n",
      "        x = F.relu(\u001b[36mself\u001b[39;49;00m.fc1(x))\r\n",
      "        x = F.dropout(x, training=\u001b[36mself\u001b[39;49;00m.training)\r\n",
      "        x = \u001b[36mself\u001b[39;49;00m.fc2(x)\r\n",
      "        \u001b[34mreturn\u001b[39;49;00m F.log_softmax(x, dim=\u001b[34m1\u001b[39;49;00m)\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32m_get_train_data_loader\u001b[39;49;00m(batch_size, training_dir, is_distributed, **kwargs):\r\n",
      "    logger.info(\u001b[33m\"\u001b[39;49;00m\u001b[33mGet train data loader\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\r\n",
      "    dataset = datasets.MNIST(training_dir, train=\u001b[36mTrue\u001b[39;49;00m, transform=transforms.Compose([\r\n",
      "        transforms.ToTensor(),\r\n",
      "        transforms.Normalize((\u001b[34m0.1307\u001b[39;49;00m,), (\u001b[34m0.3081\u001b[39;49;00m,))\r\n",
      "    ]))\r\n",
      "    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset) \u001b[34mif\u001b[39;49;00m is_distributed \u001b[34melse\u001b[39;49;00m \u001b[36mNone\u001b[39;49;00m\r\n",
      "    \u001b[34mreturn\u001b[39;49;00m torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=train_sampler \u001b[35mis\u001b[39;49;00m \u001b[36mNone\u001b[39;49;00m,\r\n",
      "                                       sampler=train_sampler, **kwargs)\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32m_get_test_data_loader\u001b[39;49;00m(test_batch_size, training_dir, **kwargs):\r\n",
      "    logger.info(\u001b[33m\"\u001b[39;49;00m\u001b[33mGet test data loader\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\r\n",
      "    \u001b[34mreturn\u001b[39;49;00m torch.utils.data.DataLoader(\r\n",
      "        datasets.MNIST(training_dir, train=\u001b[36mFalse\u001b[39;49;00m, transform=transforms.Compose([\r\n",
      "            transforms.ToTensor(),\r\n",
      "            transforms.Normalize((\u001b[34m0.1307\u001b[39;49;00m,), (\u001b[34m0.3081\u001b[39;49;00m,))\r\n",
      "        ])),\r\n",
      "        batch_size=test_batch_size, shuffle=\u001b[36mTrue\u001b[39;49;00m, **kwargs)\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32m_average_gradients\u001b[39;49;00m(model):\r\n",
      "    \u001b[37m# Gradient averaging.\u001b[39;49;00m\r\n",
      "    size = \u001b[36mfloat\u001b[39;49;00m(dist.get_world_size())\r\n",
      "    \u001b[34mfor\u001b[39;49;00m param \u001b[35min\u001b[39;49;00m model.parameters():\r\n",
      "        dist.all_reduce(param.grad.data, op=dist.reduce_op.SUM)\r\n",
      "        param.grad.data /= size\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mtrain\u001b[39;49;00m(args):\r\n",
      "    is_distributed = \u001b[36mlen\u001b[39;49;00m(args.hosts) > \u001b[34m1\u001b[39;49;00m \u001b[35mand\u001b[39;49;00m args.backend \u001b[35mis\u001b[39;49;00m \u001b[35mnot\u001b[39;49;00m \u001b[36mNone\u001b[39;49;00m\r\n",
      "    logger.debug(\u001b[33m\"\u001b[39;49;00m\u001b[33mDistributed training - {}\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m.format(is_distributed))\r\n",
      "    use_cuda = args.num_gpus > \u001b[34m0\u001b[39;49;00m\r\n",
      "    logger.debug(\u001b[33m\"\u001b[39;49;00m\u001b[33mNumber of gpus available - {}\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m.format(args.num_gpus))\r\n",
      "    kwargs = {\u001b[33m'\u001b[39;49;00m\u001b[33mnum_workers\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m: \u001b[34m1\u001b[39;49;00m, \u001b[33m'\u001b[39;49;00m\u001b[33mpin_memory\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m: \u001b[36mTrue\u001b[39;49;00m} \u001b[34mif\u001b[39;49;00m use_cuda \u001b[34melse\u001b[39;49;00m {}\r\n",
      "    device = torch.device(\u001b[33m\"\u001b[39;49;00m\u001b[33mcuda\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m \u001b[34mif\u001b[39;49;00m use_cuda \u001b[34melse\u001b[39;49;00m \u001b[33m\"\u001b[39;49;00m\u001b[33mcpu\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\r\n",
      "\r\n",
      "    \u001b[34mif\u001b[39;49;00m is_distributed:\r\n",
      "        \u001b[37m# Initialize the distributed environment.\u001b[39;49;00m\r\n",
      "        world_size = \u001b[36mlen\u001b[39;49;00m(args.hosts)\r\n",
      "        os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mWORLD_SIZE\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m] = \u001b[36mstr\u001b[39;49;00m(world_size)\r\n",
      "        host_rank = args.hosts.index(args.current_host)\r\n",
      "        os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mRANK\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m] = \u001b[36mstr\u001b[39;49;00m(host_rank)\r\n",
      "        dist.init_process_group(backend=args.backend, rank=host_rank, world_size=world_size)\r\n",
      "        logger.info(\u001b[33m'\u001b[39;49;00m\u001b[33mInitialized the distributed environment: \u001b[39;49;00m\u001b[33m\\'\u001b[39;49;00m\u001b[33m{}\u001b[39;49;00m\u001b[33m\\'\u001b[39;49;00m\u001b[33m backend on {} nodes. \u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(\r\n",
      "            args.backend, dist.get_world_size()) + \u001b[33m'\u001b[39;49;00m\u001b[33mCurrent host rank is {}. Number of gpus: {}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(\r\n",
      "            dist.get_rank(), args.num_gpus))\r\n",
      "\r\n",
      "    \u001b[37m# set the seed for generating random numbers\u001b[39;49;00m\r\n",
      "    torch.manual_seed(args.seed)\r\n",
      "    \u001b[34mif\u001b[39;49;00m use_cuda:\r\n",
      "        torch.cuda.manual_seed(args.seed)\r\n",
      "\r\n",
      "    train_loader = _get_train_data_loader(args.batch_size, args.data_dir, is_distributed, **kwargs)\r\n",
      "    test_loader = _get_test_data_loader(args.test_batch_size, args.data_dir, **kwargs)\r\n",
      "\r\n",
      "    logger.debug(\u001b[33m\"\u001b[39;49;00m\u001b[33mProcesses {}/{} ({:.0f}\u001b[39;49;00m\u001b[33m%\u001b[39;49;00m\u001b[33m) of train data\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m.format(\r\n",
      "        \u001b[36mlen\u001b[39;49;00m(train_loader.sampler), \u001b[36mlen\u001b[39;49;00m(train_loader.dataset),\r\n",
      "        \u001b[34m100.\u001b[39;49;00m * \u001b[36mlen\u001b[39;49;00m(train_loader.sampler) / \u001b[36mlen\u001b[39;49;00m(train_loader.dataset)\r\n",
      "    ))\r\n",
      "\r\n",
      "    logger.debug(\u001b[33m\"\u001b[39;49;00m\u001b[33mProcesses {}/{} ({:.0f}\u001b[39;49;00m\u001b[33m%\u001b[39;49;00m\u001b[33m) of test data\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m.format(\r\n",
      "        \u001b[36mlen\u001b[39;49;00m(test_loader.sampler), \u001b[36mlen\u001b[39;49;00m(test_loader.dataset),\r\n",
      "        \u001b[34m100.\u001b[39;49;00m * \u001b[36mlen\u001b[39;49;00m(test_loader.sampler) / \u001b[36mlen\u001b[39;49;00m(test_loader.dataset)\r\n",
      "    ))\r\n",
      "\r\n",
      "    model = Net().to(device)\r\n",
      "    \u001b[34mif\u001b[39;49;00m is_distributed \u001b[35mand\u001b[39;49;00m use_cuda:\r\n",
      "        \u001b[37m# multi-machine multi-gpu case\u001b[39;49;00m\r\n",
      "        model = torch.nn.parallel.DistributedDataParallel(model)\r\n",
      "    \u001b[34melse\u001b[39;49;00m:\r\n",
      "        \u001b[37m# single-machine multi-gpu case or single-machine or multi-machine cpu case\u001b[39;49;00m\r\n",
      "        model = torch.nn.DataParallel(model)\r\n",
      "\r\n",
      "    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)\r\n",
      "\r\n",
      "    \u001b[34mfor\u001b[39;49;00m epoch \u001b[35min\u001b[39;49;00m \u001b[36mrange\u001b[39;49;00m(\u001b[34m1\u001b[39;49;00m, args.epochs + \u001b[34m1\u001b[39;49;00m):\r\n",
      "        model.train()\r\n",
      "        \u001b[34mfor\u001b[39;49;00m batch_idx, (data, target) \u001b[35min\u001b[39;49;00m \u001b[36menumerate\u001b[39;49;00m(train_loader, \u001b[34m1\u001b[39;49;00m):\r\n",
      "            data, target = data.to(device), target.to(device)\r\n",
      "            optimizer.zero_grad()\r\n",
      "            output = model(data)\r\n",
      "            loss = F.nll_loss(output, target)\r\n",
      "            loss.backward()\r\n",
      "            \u001b[34mif\u001b[39;49;00m is_distributed \u001b[35mand\u001b[39;49;00m \u001b[35mnot\u001b[39;49;00m use_cuda:\r\n",
      "                \u001b[37m# average gradients manually for multi-machine cpu case only\u001b[39;49;00m\r\n",
      "                _average_gradients(model)\r\n",
      "            optimizer.step()\r\n",
      "            \u001b[34mif\u001b[39;49;00m batch_idx % args.log_interval == \u001b[34m0\u001b[39;49;00m:\r\n",
      "                logger.info(\u001b[33m'\u001b[39;49;00m\u001b[33mTrain Epoch: {} [{}/{} ({:.0f}\u001b[39;49;00m\u001b[33m%\u001b[39;49;00m\u001b[33m)] Loss: {:.6f}\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(\r\n",
      "                    epoch, batch_idx * \u001b[36mlen\u001b[39;49;00m(data), \u001b[36mlen\u001b[39;49;00m(train_loader.sampler),\r\n",
      "                    \u001b[34m100.\u001b[39;49;00m * batch_idx / \u001b[36mlen\u001b[39;49;00m(train_loader), loss.item()))\r\n",
      "        test(model, test_loader, device)\r\n",
      "    save_model(model, args.model_dir)\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mtest\u001b[39;49;00m(model, test_loader, device):\r\n",
      "    model.eval()\r\n",
      "    test_loss = \u001b[34m0\u001b[39;49;00m\r\n",
      "    correct = \u001b[34m0\u001b[39;49;00m\r\n",
      "    \u001b[34mwith\u001b[39;49;00m torch.no_grad():\r\n",
      "        \u001b[34mfor\u001b[39;49;00m data, target \u001b[35min\u001b[39;49;00m test_loader:\r\n",
      "            data, target = data.to(device), target.to(device)\r\n",
      "            output = model(data)\r\n",
      "            test_loss += F.nll_loss(output, target, size_average=\u001b[36mFalse\u001b[39;49;00m).item()  \u001b[37m# sum up batch loss\u001b[39;49;00m\r\n",
      "            pred = output.max(\u001b[34m1\u001b[39;49;00m, keepdim=\u001b[36mTrue\u001b[39;49;00m)[\u001b[34m1\u001b[39;49;00m]  \u001b[37m# get the index of the max log-probability\u001b[39;49;00m\r\n",
      "            correct += pred.eq(target.view_as(pred)).sum().item()\r\n",
      "\r\n",
      "    test_loss /= \u001b[36mlen\u001b[39;49;00m(test_loader.dataset)\r\n",
      "    logger.info(\u001b[33m'\u001b[39;49;00m\u001b[33mTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}\u001b[39;49;00m\u001b[33m%\u001b[39;49;00m\u001b[33m)\u001b[39;49;00m\u001b[33m\\n\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m.format(\r\n",
      "        test_loss, correct, \u001b[36mlen\u001b[39;49;00m(test_loader.dataset),\r\n",
      "        \u001b[34m100.\u001b[39;49;00m * correct / \u001b[36mlen\u001b[39;49;00m(test_loader.dataset)))\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mmodel_fn\u001b[39;49;00m(model_dir):\r\n",
      "    device = torch.device(\u001b[33m\"\u001b[39;49;00m\u001b[33mcuda\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m \u001b[34mif\u001b[39;49;00m torch.cuda.is_available() \u001b[34melse\u001b[39;49;00m \u001b[33m\"\u001b[39;49;00m\u001b[33mcpu\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\r\n",
      "    model = torch.nn.DataParallel(Net())\r\n",
      "    \u001b[34mwith\u001b[39;49;00m \u001b[36mopen\u001b[39;49;00m(os.path.join(model_dir, \u001b[33m'\u001b[39;49;00m\u001b[33mmodel.pth\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m), \u001b[33m'\u001b[39;49;00m\u001b[33mrb\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m) \u001b[34mas\u001b[39;49;00m f:\r\n",
      "        model.load_state_dict(torch.load(f))\r\n",
      "    \u001b[34mreturn\u001b[39;49;00m model.to(device)\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32msave_model\u001b[39;49;00m(model, model_dir):\r\n",
      "    logger.info(\u001b[33m\"\u001b[39;49;00m\u001b[33mSaving the model.\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\r\n",
      "    path = os.path.join(model_dir, \u001b[33m'\u001b[39;49;00m\u001b[33mmodel.pth\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    \u001b[37m# recommended way from http://pytorch.org/docs/master/notes/serialization.html\u001b[39;49;00m\r\n",
      "    torch.save(model.cpu().state_dict(), path)\r\n",
      "\r\n",
      "\r\n",
      "\u001b[34mif\u001b[39;49;00m \u001b[31m__name__\u001b[39;49;00m == \u001b[33m'\u001b[39;49;00m\u001b[33m__main__\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m:\r\n",
      "    parser = argparse.ArgumentParser()\r\n",
      "\r\n",
      "    \u001b[37m# Data and model checkpoints directories\u001b[39;49;00m\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--batch-size\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=\u001b[34m64\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mN\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33minput batch size for training (default: 64)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--test-batch-size\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=\u001b[34m1000\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mN\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33minput batch size for testing (default: 1000)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--epochs\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=\u001b[34m10\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mN\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mnumber of epochs to train (default: 10)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--lr\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mfloat\u001b[39;49;00m, default=\u001b[34m0.01\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mLR\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mlearning rate (default: 0.01)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--momentum\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mfloat\u001b[39;49;00m, default=\u001b[34m0.5\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mM\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mSGD momentum (default: 0.5)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--seed\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=\u001b[34m1\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mS\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mrandom seed (default: 1)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--log-interval\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=\u001b[34m100\u001b[39;49;00m, metavar=\u001b[33m'\u001b[39;49;00m\u001b[33mN\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mhow many batches to wait before logging training status\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--backend\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, default=\u001b[36mNone\u001b[39;49;00m,\r\n",
      "                        help=\u001b[33m'\u001b[39;49;00m\u001b[33mbackend for distributed training (tcp, gloo on cpu and gloo, nccl on gpu)\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m)\r\n",
      "\r\n",
      "    \u001b[37m# Container environment\u001b[39;49;00m\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--hosts\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mlist\u001b[39;49;00m, default=json.loads(os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_HOSTS\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m]))\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--current-host\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, default=os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_CURRENT_HOST\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--model-dir\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, default=os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_MODEL_DIR\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--data-dir\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, default=os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_CHANNEL_TRAINING\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "    parser.add_argument(\u001b[33m'\u001b[39;49;00m\u001b[33m--num-gpus\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=os.environ[\u001b[33m'\u001b[39;49;00m\u001b[33mSM_NUM_GPUS\u001b[39;49;00m\u001b[33m'\u001b[39;49;00m])\r\n",
      "\r\n",
      "    train(parser.parse_args())\r\n"
     ]
    }
   ],
   "source": [
    "!pygmentize mnist.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use Step Functions to run training in SageMaker\n",
    "\n",
    "The `PyTorch` class allows us to run our training function as a training job on SageMaker. We need to configure it with our training script, an IAM role, the number of training instances, the training instance type, and hyperparameters. In this case we are going to run our training job on 2 ```ml.c4.xlarge``` instances. But this example can be ran on one or multiple, cpu or gpu instances ([full list of available instances](https://aws.amazon.com/sagemaker/pricing/instance-types/)). The hyperparameters parameter is a dict of values that will be passed to your training script -- you can see how to access these values in the `mnist.py` script above.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.pytorch import PyTorch\n",
    "\n",
    "estimator = PyTorch(entry_point='mnist.py',\n",
    "                    role=sagemaker_execution_role,\n",
    "                    framework_version='1.2.0',\n",
    "                    train_instance_count=2,\n",
    "                    train_instance_type='ml.c4.xlarge',\n",
    "                    hyperparameters={\n",
    "                        'epochs': 6,\n",
    "                        'backend': 'gloo'\n",
    "                    })"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build a training pipeline with the Step Functions SDK\n",
    "\n",
    "A typical task for a data scientist is to train a model and deploy that model to an endpoint. Without the Step Functions SDK, this is a four step process on SageMaker that includes the following.\n",
    "\n",
    "1. Training the model\n",
    "2. Creating the model on SageMaker\n",
    "3. Creating an endpoint configuration\n",
    "4. Deploying the trained model to the configured endpoint\n",
    "\n",
    "The Step Functions SDK provides the [TrainingPipeline](https://aws-step-functions-data-science-sdk.readthedocs.io/en/latest/pipelines.html#stepfunctions.template.pipeline.train.TrainingPipeline) API to simplify this procedure. The following configures `pipeline` with the necessary parameters to define a training pipeline.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = TrainingPipeline(\n",
    "    estimator=estimator,\n",
    "    role=workflow_execution_role,\n",
    "    inputs=inputs,\n",
    "    s3_bucket=bucket\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can now view the workflow definition, and also visualize it as a graph. This workflow and graph represent your training pipeline.\n",
    "\n",
    "#### View the workflow definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"StartAt\": \"Training\",\n",
      "    \"States\": {\n",
      "        \"Training\": {\n",
      "            \"Resource\": \"arn:aws:states:::sagemaker:createTrainingJob.sync\",\n",
      "            \"Parameters\": {\n",
      "                \"AlgorithmSpecification.$\": \"$$.Execution.Input['Training'].AlgorithmSpecification\",\n",
      "                \"OutputDataConfig.$\": \"$$.Execution.Input['Training'].OutputDataConfig\",\n",
      "                \"StoppingCondition.$\": \"$$.Execution.Input['Training'].StoppingCondition\",\n",
      "                \"ResourceConfig.$\": \"$$.Execution.Input['Training'].ResourceConfig\",\n",
      "                \"RoleArn.$\": \"$$.Execution.Input['Training'].RoleArn\",\n",
      "                \"InputDataConfig.$\": \"$$.Execution.Input['Training'].InputDataConfig\",\n",
      "                \"HyperParameters.$\": \"$$.Execution.Input['Training'].HyperParameters\",\n",
      "                \"TrainingJobName.$\": \"$$.Execution.Input['Training'].TrainingJobName\"\n",
      "            },\n",
      "            \"Type\": \"Task\",\n",
      "            \"Next\": \"Create Model\"\n",
      "        },\n",
      "        \"Create Model\": {\n",
      "            \"Parameters\": {\n",
      "                \"ModelName.$\": \"$$.Execution.Input['Create Model'].ModelName\",\n",
      "                \"PrimaryContainer.$\": \"$$.Execution.Input['Create Model'].PrimaryContainer\",\n",
      "                \"ExecutionRoleArn.$\": \"$$.Execution.Input['Create Model'].ExecutionRoleArn\"\n",
      "            },\n",
      "            \"Resource\": \"arn:aws:states:::sagemaker:createModel\",\n",
      "            \"Type\": \"Task\",\n",
      "            \"Next\": \"Configure Endpoint\"\n",
      "        },\n",
      "        \"Configure Endpoint\": {\n",
      "            \"Resource\": \"arn:aws:states:::sagemaker:createEndpointConfig\",\n",
      "            \"Parameters\": {\n",
      "                \"EndpointConfigName.$\": \"$$.Execution.Input['Configure Endpoint'].EndpointConfigName\",\n",
      "                \"ProductionVariants.$\": \"$$.Execution.Input['Configure Endpoint'].ProductionVariants\"\n",
      "            },\n",
      "            \"Type\": \"Task\",\n",
      "            \"Next\": \"Deploy\"\n",
      "        },\n",
      "        \"Deploy\": {\n",
      "            \"Resource\": \"arn:aws:states:::sagemaker:createEndpoint\",\n",
      "            \"Parameters\": {\n",
      "                \"EndpointConfigName.$\": \"$$.Execution.Input['Deploy'].EndpointConfigName\",\n",
      "                \"EndpointName.$\": \"$$.Execution.Input['Deploy'].EndpointName\"\n",
      "            },\n",
      "            \"Type\": \"Task\",\n",
      "            \"End\": true\n",
      "        }\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "print(pipeline.workflow.definition.to_json(pretty=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Visualize the workflow graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<link rel=\"stylesheet\" type=\"text/css\" href=\"https://do0of8uwbahzz.cloudfront.net/graph.css\">\n",
       "<div id=\"graph-941\" class=\"workflowgraph\">\n",
       "    \n",
       "    <svg></svg>\n",
       "    \n",
       "</div>\n",
       "\n",
       "<script type=\"text/javascript\">\n",
       "\n",
       "require.config({\n",
       "    paths: {\n",
       "        sfn: \"https://do0of8uwbahzz.cloudfront.net/sfn\",\n",
       "    }\n",
       "});\n",
       "\n",
       "require(['sfn'], function(sfn) {\n",
       "    var options = {\n",
       "        width: $('#graph-941').width(),\n",
       "        height: 600,\n",
       "        layout: 'LR',\n",
       "        resizeHeight: true\n",
       "    };\n",
       "\n",
       "    var definition = {\"StartAt\": \"Training\", \"States\": {\"Training\": {\"Resource\": \"arn:aws:states:::sagemaker:createTrainingJob.sync\", \"Parameters\": {\"AlgorithmSpecification.$\": \"$$.Execution.Input['Training'].AlgorithmSpecification\", \"OutputDataConfig.$\": \"$$.Execution.Input['Training'].OutputDataConfig\", \"StoppingCondition.$\": \"$$.Execution.Input['Training'].StoppingCondition\", \"ResourceConfig.$\": \"$$.Execution.Input['Training'].ResourceConfig\", \"RoleArn.$\": \"$$.Execution.Input['Training'].RoleArn\", \"InputDataConfig.$\": \"$$.Execution.Input['Training'].InputDataConfig\", \"HyperParameters.$\": \"$$.Execution.Input['Training'].HyperParameters\", \"TrainingJobName.$\": \"$$.Execution.Input['Training'].TrainingJobName\"}, \"Type\": \"Task\", \"Next\": \"Create Model\"}, \"Create Model\": {\"Parameters\": {\"ModelName.$\": \"$$.Execution.Input['Create Model'].ModelName\", \"PrimaryContainer.$\": \"$$.Execution.Input['Create Model'].PrimaryContainer\", \"ExecutionRoleArn.$\": \"$$.Execution.Input['Create Model'].ExecutionRoleArn\"}, \"Resource\": \"arn:aws:states:::sagemaker:createModel\", \"Type\": \"Task\", \"Next\": \"Configure Endpoint\"}, \"Configure Endpoint\": {\"Resource\": \"arn:aws:states:::sagemaker:createEndpointConfig\", \"Parameters\": {\"EndpointConfigName.$\": \"$$.Execution.Input['Configure Endpoint'].EndpointConfigName\", \"ProductionVariants.$\": \"$$.Execution.Input['Configure Endpoint'].ProductionVariants\"}, \"Type\": \"Task\", \"Next\": \"Deploy\"}, \"Deploy\": {\"Resource\": \"arn:aws:states:::sagemaker:createEndpoint\", \"Parameters\": {\"EndpointConfigName.$\": \"$$.Execution.Input['Deploy'].EndpointConfigName\", \"EndpointName.$\": \"$$.Execution.Input['Deploy'].EndpointName\"}, \"Type\": \"Task\", \"End\": true}}};\n",
       "    var elementId = '#graph-941';\n",
       "\n",
       "    var graph = new sfn.StateMachineGraph(definition, elementId, options);\n",
       "    graph.render();\n",
       "});\n",
       "\n",
       "</script>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipeline.render_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and execute the pipeline on AWS Step Functions\n",
    "\n",
    "Create the pipeline in AWS Step Functions with [create](https://aws-step-functions-data-science-sdk.readthedocs.io/en/latest/workflow.html#stepfunctions.workflow.Workflow.create)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32m[INFO] Workflow created successfully on AWS Step Functions.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'arn:aws:states:us-east-1:806570384721:stateMachine:training-pipeline-2020-01-14-16-42-55'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipeline.create()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the workflow with [execute](https://aws-step-functions-data-science-sdk.readthedocs.io/en/latest/workflow.html#stepfunctions.workflow.Workflow.execute). A link will be provided after the following cell is executed. Following this link, you can monitor your pipeline execution on Step Functions' console."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32m[INFO] Workflow execution started successfully on AWS Step Functions.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Execution: <a target=\"_blank\" href=\"https://console.aws.amazon.com/states/home?region=us-east-1#/executions/details/arn:aws:states:us-east-1:806570384721:execution:training-pipeline-2020-01-14-16-42-55:training-pipeline-2020-01-14-16-43-07\">arn:aws:states:us-east-1:806570384721:execution:training-pipeline-2020-01-14-16-42-55:training-pipeline-2020-01-14-16-43-07</a>"
      ],
      "text/plain": [
       "Execution(execution_arn='arn:aws:states:us-east-1:806570384721:execution:training-pipeline-2020-01-14-16-42-55:training-pipeline-2020-01-14-16-43-07', name=None, status=<ExecutionStatus.Running: 'RUNNING'>, start_date=datetime.datetime(2020, 1, 14, 16, 43, 7, 364000, tzinfo=tzlocal()))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pipeline.execute()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p36",
   "language": "python",
   "name": "conda_pytorch_p36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
